# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import math
import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

FLAGS = flags.FLAGS


def do_plot_spread(df, csv_root, exchange, market):
  fig, ax1 = plt.subplots(figsize=(20, 10))

  ax2 = ax1.twinx()
  ax1.tick_params(axis='x', rotation=45)
  symbols = df['symbol']
  plt.title('aggregated bid_ask_spread bp VS volume in %s %s market 20180725-20180730' %
            (exchange, market))
  plt.xticks(df.index, symbols)
  ax1.plot(df.index, df['avg_bid_ask_spread/vwap_bp'], color='g', marker='o')
  ax2.plot(df.index, df['total_volume*vwap'], color='b', marker='o')
  tick_space = (ax1.get_yticks()[-1] - ax1.get_yticks()[0]) / 20
  tick_space = int(math.ceil(tick_space / 5)) * 5
  ax1.yaxis.set_major_locator(ticker.MultipleLocator(tick_space))

  ax1.set_ylabel('bid_ask_spread bp', color='g')
  ax1.grid(True)
  ax2.set_ylabel('volume %s' % market, color='b')
  plt.savefig(os.path.join(csv_root, 'aggregated_%s_%s_20180725-20180730.png' % (exchange, market)))


def aggregate_spread_volume(dfs, exchange, market):
  extracted_dfs = []
  for df in dfs:
    extracted_df = df.loc[(df['exchange'] == exchange) & (df['market'] == market)].head(100)
    extracted_dfs.append(extracted_df)

  concatted_df = pandas.DataFrame()
  for extracted_df in extracted_dfs:
    concatted_df = pandas.concat([concatted_df, extracted_df])

  agg_df = concatted_df.groupby('symbol')['avg_bid_ask_spread/vwap_bp', 'total_volume*vwap'].mean()
  agg_df = agg_df.reset_index().sort_values('total_volume*vwap',
                                            ascending=False).reset_index().head(30)
  return agg_df


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  assert csv_root, '--csv_root must be specified.'

  spread_dfs = []
  for date in range(20180725, 20180730, 1):
    spread_df = read_csv_into_df(csv_root, 'everything_%d.csv.normalized' % date)
    spread_dfs.append(spread_df)

  for exchange in ('Huobi',):
    for market in ('BTC',):
      aggregated_df = aggregate_spread_volume(spread_dfs, exchange, market)
      do_plot_spread(aggregated_df, csv_root, exchange, market)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv files root directory.')

  app.run(main)
